#!/bin/bash

TOP_DIR="$(dirname "${BASH_SOURCE[0]}")"

# Default settings
DEFAULT_PROGRAMS=("chss_pkg" "slam_node" "pnc_node" "costmap_node")
DEFAULT_CSV_DIR="${TOP_DIR}/csv"
DEFAULT_OUTPUT_DIR="${TOP_DIR}/out"
DEFAULT_DURATION=86400  # 24 hours in seconds
PYTHON_XLSX_PROC="${TOP_DIR}/scripts/xlxs_process.py"
PYTHON_CHARTS_PROC="${TOP_DIR}/scripts/generate_charts_from_raw.py"

# Function to display usage information
usage() {
    echo "Usage: $0 [OPTIONS]"
    echo "Options:"
    echo "  -r, --run PROGRAM1 [PROGRAM2 ...]  Run resource analysis for specified programs"
    echo "  -p, --process                      Process collected data and generate Excel report"
    echo "  -d, --duration SECONDS             Set the duration for data collection (default: 86400 seconds)"
    echo "  -h, --help                         Display this help message"
    echo
    echo "Examples:"
    echo "  $0 -r program1 program2            Run analysis for program1 and program2"
    echo "  $0 -r                              Run analysis for default programs"
    echo "  $0 -p                              Process data and generate Excel report"
    echo "  $0 -r program1 -d 3600             Run analysis for program1 for 1 hour"
}

# Function to set up CSV files
setup_csv_files() {
    local programs=("$@")
    mkdir -p $DEFAULT_CSV_DIR
    rm -f $DEFAULT_CSV_DIR/*.csv

    testing_time=$(date +"%Y-%m-%d %H:%M:%S")
    for program in "${programs[@]}"; do
        output_file="$DEFAULT_CSV_DIR/${program}_resource_usage.csv"
        echo -e "TESTING DATE: $testing_time \n" > "$output_file"
        echo "TIMESTAMP,PROGRAM,PID,%MEM,%CPU,VSZ,RSS,TIME" >> "$output_file"
    done

    echo -e "TESTING DATE: $testing_time \n" > "$DEFAULT_CSV_DIR/cpu_usage.csv"
    echo "TIME,CPU,%USR,%NICE,%SYS,%IOWAIT,%IRQ,%SOFT,%STEAL,%GUEST,%GNICE,%IDLE,TEMP" >> "$DEFAULT_CSV_DIR/cpu_usage.csv"

    echo -e "TESTING DATE: $testing_time \n" > "$DEFAULT_CSV_DIR/mem_usage.csv"
    echo "TIME,KBMEMFREE,KBAVAIL,KBMEMUSED,%MEMUSED,KBBUFFERS,KBCACHED,KBCOMMIT,%COMMIT,KBACTIVE,KBINACT,KBDIRTY" >> "$DEFAULT_CSV_DIR/mem_usage.csv"
}

# Function to collect CPU data
cpu_data_collection() {
    LC_TIME=C mpstat -P ALL 1 1 | awk -v OFS=',' 'NR>3&&NR<13 {
        print $1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12
    }' >> "$DEFAULT_CSV_DIR/cpu_usage.csv"
}

# Function to collect memory data
mem_data_collection() {
    sar -r 1 1 | awk -v OFS=',' 'NR==4 {
        gsub(/AM|PM/,"",$2);
        printf "%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s\n",
        $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13
    }' >> "$DEFAULT_CSV_DIR/mem_usage.csv"
}

# Function to run data collection
run_data_collection() {
    local programs=("$@")
    local duration=$DEFAULT_DURATION
    local end_time=$((SECONDS + duration))
    local interrupt_requested=false

    setup_csv_files "${programs[@]}"
    echo "Starting resource analysis for programs: ${programs[*]}"
    echo "Data collection will run for $duration seconds. Press Ctrl+C to stop early."

    trap 'echo "Interrupt received. Finishing current data collection cycle..."; interrupt_requested=true' INT

    while [ $SECONDS -lt $end_time ] && ! $interrupt_requested; do
        for p in "${programs[@]}"; do
            output_file="$DEFAULT_CSV_DIR/${p}_resource_usage.csv"
            if pid=$(pgrep -o "$p"); then
                ps -p "$pid" -o pid=,%mem=,%cpu=,vsz=,rss=,time= --no-headers |
                awk -v time="$(date +"%H:%M:%S")" -v prog="$p" '{printf "%s,%s,%s,%s,%s,%s,%s,%s\n", time, prog, $1, $2, $3, $4, $5, $6}' >> "$output_file"
            else
                echo "$(date +"%H:%M:%S"),$p,Not Running,0,0,0,0,0" >> "$output_file"
            fi
        done

        cpu_data_collection
        mem_data_collection

        # Calculate and display remaining time
        remaining_time=$((end_time - SECONDS))
        printf "\rRemaining time: %02d:%02d:%02d" $((remaining_time/3600)) $((remaining_time%3600/60)) $((remaining_time%60))
    done

    echo -e "\nData collection completed."
    trap - INT
}

# Function to process data and create Excel file
process_data() {
    echo "Processing data and creating raw Excel file..."
    if [ ! -d "$DEFAULT_CSV_DIR" ] || [ -z "$(ls -A $DEFAULT_CSV_DIR)" ]; then
        echo "Error: No data found in $DEFAULT_CSV_DIR. Please run data collection first."
        exit 1
    fi

    python3 "$PYTHON_XLSX_PROC" "$DEFAULT_CSV_DIR" "$DEFAULT_OUTPUT_DIR"
    if [ $? -ne 0 ]; then
        echo "Error occurred while creating raw Excel file"
        exit 1
    fi

    echo "Creating analysis charts..."
    latest_raw_file=$(ls -t "$DEFAULT_OUTPUT_DIR"/resource_raw_*.xlsx | head -n1)
    if [ ! -f "$PYTHON_CHARTS_PROC" ]; then
        echo "Error: Analysis script not found at $PYTHON_CHARTS_PROC"
        exit 1
    fi
    python3 "$PYTHON_CHARTS_PROC" "$latest_raw_file" "$DEFAULT_OUTPUT_DIR"
    if [ $? -eq 0 ]; then
        echo "Analysis charts created successfully in $DEFAULT_OUTPUT_DIR"
    else
        echo "Error occurred while creating analysis charts"
        exit 1
    fi
}

# Main execution
if [ $# -eq 0 ]; then
    usage
    exit 1
fi

MODE=""
PROGRAMS=()

while [[ $# -gt 0 ]]; do
    case $1 in
        -r|--run)
            MODE="run"
            shift
            while [[ $# -gt 0 ]] && [[ "$1" != -* ]]; do
                PROGRAMS+=("$1")
                shift
            done
            [ ${#PROGRAMS[@]} -eq 0 ] && PROGRAMS=("${DEFAULT_PROGRAMS[@]}")
            ;;
        -p|--process)
            MODE="process"
            shift
            ;;
        -d|--duration)
            DEFAULT_DURATION=$2
            shift 2
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            usage
            exit 1
            ;;
    esac
done

case $MODE in
    "run")
        run_data_collection "${PROGRAMS[@]}"
        ;;
    "process")
        process_data
        ;;
    *)
        echo "Error: No valid mode specified."
        usage
        exit 1
        ;;
esac
